import numpy as np

from experimental.prophet import op_registry
from experimental.prophet.kernels.op_kernel import OpKernel, arg_def
from experimental.prophet.graph import graph
from experimental.prophet.ops.constant import constant


class TestPyOp(OpKernel):
  name = 'TestPyOp'
  input_args = [arg_def('in', float)]
  output_args = [arg_def('out', float)]

  def __init__(self, attrs, node_def):
    print('init!')
    print(node_def)
    print('-----')

  def compute(self, inputs, outputs):
    v = inputs[0].value
    outputs[0].value = np.matmul(v, v)
    print(outputs[0].value)


op_registry.register_op(TestPyOp)

g = graph.get_default_graph()
c = constant([[1, 3], [5., 7]])
n = graph.create_node('TestPyOp', [c])

g.eval()
print(n.outputs[0].value)
